from flask import Flask, render_template, Response
import zmq
import numpy as np
import cv2
import threading
import json
import signal
import time
import os

app = Flask(__name__)

zmq_host = os.environ.get('ZMQ_HOST', 'default_host')
address = f"tcp://{zmq_host}:5555"
topic = "hello"

# ZeroMQ context and socket
context = zmq.Context()
socket = context.socket(zmq.SUB)
socket.connect(address)
socket.setsockopt_string(zmq.SUBSCRIBE, topic)
print("Waiting messages...")

# A dictionary to store images, using a unique ID as key
image_queue = []
stop_thread = False

# Function to receive images via ZeroMQ
def receive_images():
    global stop_thread
    while not stop_thread:
        try:
            topic, frame_bytes, result_bytes = socket.recv_multipart()
            
            # parse topic
            topic = topic.decode()
            # parse frame
            buffer = np.frombuffer(frame_bytes, dtype=np.uint8)
            frame = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
            # parse result
            result = json.loads(result_bytes.decode('utf-8'))

            # Draw results on frame
            if result != {}:
                lines = result["lines"] if "lines" in result else []
                circles = result["circles"] if "circles" in result else []

                # Draw 2D
                for line in lines:
                    cv2.line(frame, tuple(line[0]), tuple(line[1]), (255,255,0), 4, cv2.LINE_AA,)
                for circle in circles:
                    cv2.circle(frame, tuple(circle), 3, (0, 255, 255), -1, cv2.LINE_AA,)

            # put frame int queue
            image_queue.append(frame)

            print(f"Got msg")

        except Exception as e:
            print(f"Error receiving images: {e}")
            break

# Start to thread to receive images
thread = threading.Thread(target=receive_images)
thread.daemon = True
thread.start()

# Route to display the last received image
@app.route('/')
def index():
    return render_template('index.html')

# Route to return the image file as a stream
@app.route('/image_feed')
def image_feed():
    def generate():
        while not stop_thread:
            if image_queue:
                frame = image_queue.pop(0)
                _, buffer = cv2.imencode('.jpg', frame)
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')
            else:
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + b' ' + b'\r\n')  # Send a blank frame
            time.sleep(0.0333)  # Adjust the sleep time based on the desired frame rate
    return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')
    

def singal_handler(sig, frame):
    global stop_thread
    print('Exiting...')
    stop_thread = True
    socket.close()
    thread.join()
    appctx = app.app_context()
    appctx.push()
    try:
        # Close the Flask application
        func = request.environ.get('werkzeug.server.shutdown')
        if func is None:
            raise RuntimeError('Not running with the Werkzeug Server')
        func()
    finally:
        appctx.pop()

if __name__ == '__main__':
    signal.signal(signal.SIGINT, singal_handler)
    app.run(host='0.0.0.0', port=5000)